stage=1
stop_stage=3

manifest_dir=/home/node27_tmpdata/xlgeng/pachong_10W_data/fairseq_data/manifest/chat_1000
mkdir -p $manifest_dir
raw_data_dir=/home/node27_tmpdata/xlgeng/pachong_10W_data/raw_wav/chat_1000

exp_dir=/home/node27_tmpdata/xlgeng/pachong_10W_data/fairseq_data/hubert_feat/chat_1000
split=valid
all_data=( ${split} )
ext="wav"
exp_name=hubert_large_by_xlgeng
extract_layer=9
nj=40
num_gpu=8
hubert_ckpt=/home/work_nfs15/asr_data/ckpt/origin_chinese_hubert/chinese_hubert_large.pt
feat=${exp_name}_extract_layer_${extract_layer}
feat_dir=${exp_dir}/feature/${feat}
mkdir -p $feat_dir
n_cluster=500
km_dir=${exp_dir}/k-means/${feat}
mkdir -p $km_dir

if [ ${stage} -le 0 ] && [ ${stop_stage} -ge 0 ]; then
  echo "stage 0: 得到tsv文件 from raw dir"
  valid=0.0
  python tools/wav2vec_manifest.py $raw_data_dir \
    --dest $manifest_dir \
    --ext $ext \
    --valid-percent $valid \
    --split $split
  #  --part-tsv-path /home/node22_tmpdata/xlgeng/pachong_10W_data/fairseq_data/manifest/shenghuo_1/train_small.tsv \
  #  --part-tsv-path2 /home/node22_tmpdata/xlgeng/pachong_10W_data/fairseq_data/manifest/shenghuo_1/train_medium1.tsv
fi

if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
  echo "stage 1: 提取hubert表征"
  echo "Compute Hubert (iter2) features started @ `date`"
  start=$(date +%s)
  for set in "${all_data[@]}"; do
      echo "drop feat: $set"
      for ((i = 0; i < $nj; ++i)); do
      {
          result=$((i % num_gpu))
          export CUDA_VISIBLE_DEVICES=$result
          # set 是tsv文件的文件名字； max_chunk单位是采样点是对于长音频，每次提取多少采样点，然后在维度拼接
          python tools/dump_hubert_feature.py   \
              ${manifest_dir} ${set}                        \
              $hubert_ckpt   \
              ${extract_layer} ${nj} $i ${feat_dir}/${set}  \
              --max_chunk 16000000
      }   &
      done
      wait
  done
  echo "Compute Hubert (iter2) features Done @ `date`"
  end=$(date +%s)
  echo "Execution time: $((end - start)) seconds"
fi

if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
  echo "stage 1: 学习kmeans"
  kmean_model_path_here="${km_dir}/model.mdl"
  kmain_model_path="/home/node27_tmpdata/xlgeng/pachong_10W_data/fairseq_data/hubert_feat/ximalaya_redian_2T/k-means/hubert_large_by_xlgeng_extract_layer_9/model.mdl"
  if [ -f $kmain_model_path ]; then
      echo "stage 12: Skipped - ${km_dir}/model.mdl exist @ `date`"
      cp $kmain_model_path $kmean_model_path_here
  else
      echo "stage 12: K-means clustering on Hubert (iter2) features started @ `date`"
      python ${base_dir}/simple_kmeans/learn_kmeans.py \
          ${feat_dir}/${split} ${split} ${nj} \
          $kmain_model_path ${n_cluster} \
          --percent 0.99 --batch_size 10000
      echo "stage 12: Done @ `date`"
  fi
fi

if [ ${stage} -le 3 ] && [ ${stop_stage} -ge 3 ]; then
  echo "stage 1: 得到km 标签"
  echo "stage 13: K-means application started @ `date`"
  for set in "${all_data[@]}"; do
      echo "drop label: $set"
      for ((i = 0; i < $nj; ++i)); do
      {
          result=$((i % num_gpu))
          export CUDA_VISIBLE_DEVICES=$result
          python tools/dump_km_label.py         \
              ${feat_dir}/${set} ${set} ${km_dir}/model.mdl               \
              ${nj} $i ${km_dir}/${set}_label
      } &
      done
      wait
      # merge labels for different shards
      for rank in $(seq 0 $((nj - 1))); do
          cat ${km_dir}/${set}_label/${set}_${rank}_${nj}.km
      done > ${km_dir}/${set}.km
  done
  echo "stage 13: Done @ `date`"
fi

if [ ${stage} -le 4 ] && [ ${stop_stage} -ge 4 ]; then
  echo "stage 1: 生成词典"
  echo "stage 14: Generate dict started @ `date`"
        # fairseq-preprocess                              \
        python ../../../fairseq_cli/preprocess.py                 \
            --only-source                               \
            --source-lang km                            \
            --dict-only                                 \
            --trainpref ${km_dir}/${split}                 \
            --destdir ${km_dir}                         \
            --workers 10
    echo "stage 14: Done @ `date`"
fi